from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from camera_thread import CameraThread
from PyQt5.QtCore import QTimer  # 确保导入 QTimer
import torch
import sys
import json
import os
import numpy as np
import cv2
import datetime
from matplotlib.figure import Figure
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置默认字体为 SimHei（黑体）
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


class YOLOv5CameraUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.available_classes = []  # 初始化为空列表，后续从模型中加载
        # 初始化 fullscreen_action
        self.fullscreen_action = QAction("全屏", self)  # 将 fullscreen_action 定义为类的成员变量
        self.init_ui()
        self.history_data = {}  # 用于存储历史统计数据
        self.start_time = QDateTime.currentDateTime()  # 记录程序启动时间
        self.timer = QTimer(self)  # 初始化定时器
        self.timer.timeout.connect(self.update_line_chart)  # 绑定定时器回调函数
        # 创建模型保存目录
        self.model_dir = "./models"
        os.makedirs(self.model_dir, exist_ok=True)  # 如果目录不存在则创建
        self.load_model()
        self.camera_thread = None
        self.is_fullscreen = False
        self.roi_mask = None  # 检测区域掩码
        self.drawing = False  # 是否正在绘制
        self.points = []      # 多边形顶点
        # 获取当前时间戳
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        self.output_video_path = f"./runs/videos/{timestamp}.mp4"  # 默认输出路径
        # 创建保存目录
        self.save_dir = "./runs/images"
        os.makedirs(self.save_dir, exist_ok=True)  # 如果目录不存在则创建
        
        self.is_drawing_roi = False  # 是否正在绘制 ROI
        # 历史 IP 摄像头地址
        self.history_ip_addresses = []  # 确保在这里初始化
        self.history_file = "history_ip_addresses.json"  # 保存历史记录的文件
        self.load_history_ip_addresses()  # 加载历史记录
        # 在 UI 初始化完成后，设置 IP 摄像头输入框的自动补全功能
        self.update_ip_camera_input_completer()
        self.source_path = None

    def init_ui(self):
        self.setWindowTitle("YOLOv5 实时检测工具")
        self.setGeometry(100, 100, 1280, 720)
        # 设置窗口图标
        self.setWindowIcon(QIcon("./icons/1.png"))  # 将 "icon.png" 替换为你的图标文件路径

        # 使用 QSplitter 实现可拖动调节左右区域占比
        splitter = QSplitter(Qt.Horizontal)
        self.setCentralWidget(splitter)

        # 左侧功能区 ==============================================
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)  # 正确：传递父控件
        left_layout.setSpacing(20)  # 设置控件间距
        left_layout.setContentsMargins(10, 10, 10, 10)  # 设置边距

        # 检测设置分组
        settings_group = QGroupBox()
        settings_group.setFont(QFont("Arial", 10))  # 设置字体为 Arial，大小为 12
        settings_layout = QVBoxLayout()  # 正确：无参数构造函数
        settings_layout.setContentsMargins(10, 10, 10, 30)  # 设置边距（左、上、右、下）
        settings_group.setLayout(settings_layout)

        # 添加“检测设置”标题
        settings_title = QLabel("检测设置")
        settings_title.setFont(QFont("Arial", 12, QFont.Bold))  # 设置字体为 Arial，大小为 12，加粗
        settings_title.setAlignment(Qt.AlignLeft)  # 居中对齐
        settings_layout.addWidget(settings_title)  # 将标题添加到布局中
        
        # 检测源选择
        source_group = QWidget()  # 创建一个容器
        source_layout = QHBoxLayout(source_group)  # 正确：传递父控件
        source_layout.setContentsMargins(0, 0, 0, 0)  # 去除边距

        source_layout.addWidget(QLabel("选择检测源:"), stretch=1)
        self.source_selector = QComboBox()
        self.source_selector.addItems(["本地摄像头", "IP摄像头", "视频文件", "图片文件"])
        self.source_selector.currentTextChanged.connect(self.update_source_ui)
        source_layout.addWidget(self.source_selector, stretch=8)
        settings_layout.addWidget(source_group)


        # 添加本地模型选择按钮
        self.btn_load_model = QPushButton("加载本地模型", self)
        self.btn_load_model.clicked.connect(self.load_local_model)
        settings_layout.addWidget(self.btn_load_model)

        # 添加保存路径选择按钮
        self.btn_save_path = QPushButton("选择保存路径", self)
        self.btn_save_path.clicked.connect(self.select_save_path)
        settings_layout.addWidget(self.btn_save_path)

        # IP摄像头地址输入
        self.ip_camera_label = QLabel("IP摄像头地址:")
        self.ip_camera_input = QLineEdit()
        self.ip_camera_input.setPlaceholderText("例如：rtsp://username:password@ip:port")
        settings_layout.addWidget(self.ip_camera_label)
        settings_layout.addWidget(self.ip_camera_input)
        self.ip_camera_label.setVisible(False)
        self.ip_camera_input.setVisible(False)
        
        # 添加清除历史 IP 地址的按钮
        self.btn_clear_history = QPushButton("清除历史 IP 地址", self)
        self.btn_clear_history.clicked.connect(self.clear_history_ip_addresses)
        settings_layout.addWidget(self.btn_clear_history)
        self.btn_clear_history.setVisible(False)  # 默认隐藏
        
        # 文件选择按钮
        self.file_button = QPushButton("选择文件", self)
        self.file_button.clicked.connect(self.select_file)
        settings_layout.addWidget(self.file_button)
        self.file_button.setVisible(False)

        # 置信度调节
        conf_group = QWidget()  # 创建一个容器
        conf_layout = QHBoxLayout(conf_group)  # 正确：传递父控件
        conf_layout.setContentsMargins(0, 0, 0, 0)  # 去除边距
        

        # 置信度滑块
        self.conf_slider = QSlider(Qt.Horizontal)
        self.conf_slider.setRange(0, 100)
        self.conf_slider.setValue(50)
        self.conf_slider.valueChanged.connect(self.update_conf_thres_from_slider)
        conf_layout.addWidget(self.conf_slider, stretch=5)

        # 置信度输入框
        self.conf_input = QLineEdit()
        self.conf_input.setPlaceholderText("0~1")
        self.conf_input.setValidator(QDoubleValidator(0.0, 1.0, 2))
        self.conf_input.editingFinished.connect(self.update_conf_thres_from_input)
        conf_layout.addWidget(self.conf_input, stretch=1)

        settings_layout.addWidget(QLabel("置信度阈值:"))  # 添加标签
        settings_layout.addWidget(conf_group)  # 将容器添加到设置布局中
        

        # IoU 阈值调节
        iou_group = QWidget()  # 创建一个容器
        iou_layout = QHBoxLayout(iou_group)  # 正确：传递父控件
        iou_layout.setContentsMargins(0, 0, 0, 0)  # 去除边距
        

        # IoU 滑块
        self.iou_slider = QSlider(Qt.Horizontal)
        self.iou_slider.setRange(0, 100)
        self.iou_slider.setValue(50)
        self.iou_slider.valueChanged.connect(self.update_iou_thres_from_slider)
        iou_layout.addWidget(self.iou_slider, stretch=5)

        # IoU 输入框
        self.iou_input = QLineEdit()
        self.iou_input.setPlaceholderText("0~1")
        self.iou_input.setValidator(QDoubleValidator(0.0, 1.0, 2))
        self.iou_input.editingFinished.connect(self.update_iou_thres_from_input)
        iou_layout.addWidget(self.iou_input, stretch=1)

        settings_layout.addWidget(QLabel("IoU阈值:"))  # 添加标签
        settings_layout.addWidget(iou_group)  # 将容器添加到设置布局中
        

        settings_group.setLayout(settings_layout)
        left_layout.addWidget(settings_group)

        # 类别选择分组
        class_group = QGroupBox()
        class_layout = QVBoxLayout()
        class_layout.setContentsMargins(10, 10, 10, 30)  # 设置边距（左、上、右、下）

        # 添加“选择检测类别”标题
        class_title = QLabel("选择检测类别")
        class_title.setFont(QFont("Arial", 12, QFont.Bold))  # 设置字体为 Arial，大小为 12，加粗
        class_title.setAlignment(Qt.AlignLeft)  # 居中对齐
        class_layout.addWidget(class_title)  # 将标题添加到布局中

        # 类别选择列表
        self.class_list = QListWidget()
        self.class_list.setSelectionMode(QAbstractItemView.MultiSelection)
        for cls in self.available_classes:
            self.class_list.addItem(cls)
        self.class_list.itemSelectionChanged.connect(self.update_selected_classes)  # 连接选择事件    
        class_layout.addWidget(self.class_list)

        # 添加“全选”和“取消”按钮
        button_layout = QHBoxLayout()
        self.btn_select_all = QPushButton("全选", self)
        self.btn_select_all.clicked.connect(self.select_all_classes)
        self.btn_deselect_all = QPushButton("取消", self)
        self.btn_deselect_all.clicked.connect(self.deselect_all_classes)
        button_layout.addWidget(self.btn_select_all)
        button_layout.addWidget(self.btn_deselect_all)
        class_layout.addLayout(button_layout)

        class_group.setLayout(class_layout)
        left_layout.addWidget(class_group)

        # 检测结果表格
        results_title = QLabel("检测结果")
        results_title.setFont(QFont("Arial", 12, QFont.Bold))  # 设置字体为 Arial，大小为 12，加粗
        results_title.setAlignment(Qt.AlignLeft)  # 居中对齐
        left_layout.addWidget(results_title)  # 将标题添加到布局中

        self.results_table = QTableWidget()
        self.results_table.setColumnCount(3)
        self.results_table.setHorizontalHeaderLabels(["类别", "置信度", "数量"])
        self.results_table.setEditTriggers(QTableWidget.NoEditTriggers)
        self.results_table.setSortingEnabled(True)
        self.results_table.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        # 隐藏垂直表头（去掉左边索引）
        self.results_table.verticalHeader().setVisible(False)
        left_layout.addWidget(self.results_table)

        # 设置左侧区域的最小宽度
        left_panel.setMinimumWidth(800)  # 左侧区域最小宽度为400像素
        left_panel.setMaximumWidth(800)
        left_panel.setMinimumHeight(1600)

        # 右侧显示区 ==============================================
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)

        # 视频显示区域
        self.video_label = QLabel(self)
        self.video_label.setAlignment(Qt.AlignCenter)
        self.video_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.video_label.mousePressEvent = self.mouse_press_event
        self.video_label.mouseMoveEvent = self.mouse_move_event
        self.video_label.mouseReleaseEvent = self.mouse_release_event
        self.fps_label = QLabel("FPS: 0.0")
        right_layout.addWidget(self.fps_label)
        right_layout.addWidget(self.video_label)

        # 按钮布局
        button_panel = QHBoxLayout()
        button_panel.setSpacing(10)  # 设置按钮之间的间距
        button_panel.setAlignment(Qt.AlignCenter)  # 居中对齐

        # 开始/停止按钮
        self.btn_start = QPushButton("开始检测", self)
        self.btn_start.clicked.connect(self.start_camera)
        self.btn_start.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)  # 宽度自适应，高度固定
        self.btn_start.setFixedHeight(100)  # 固定高度为100像素

        self.btn_stop = QPushButton("停止", self)
        self.btn_stop.clicked.connect(self.stop_camera)
        self.btn_stop.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
        self.btn_stop.setFixedHeight(100)
        self.btn_stop.setEnabled(False)

        # 清除检测区域按钮
        self.btn_clear_roi = QPushButton("清除检测区域", self)
        self.btn_clear_roi.clicked.connect(self.clear_roi)
        self.btn_clear_roi.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
        self.btn_clear_roi.setFixedHeight(100)

        # 绘制 ROI 按钮
        self.btn_draw_roi = QPushButton("开始绘制 ROI", self)
        self.btn_draw_roi.clicked.connect(self.toggle_draw_roi)
        self.btn_draw_roi.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Fixed)
        self.btn_draw_roi.setFixedHeight(100)

        # 创建菜单栏
        self.menu_bar = self.menuBar()

        # 添加“统计工具”菜单
        stats_menu = self.menu_bar.addMenu("统计工具")

        # 添加“折线图”菜单项
        line_chart_action = QAction("折线图", self)
        line_chart_action.triggered.connect(self.show_line_chart)
        stats_menu.addAction(line_chart_action)

        # 添加“官方模型”菜单
        model_menu = self.menu_bar.addMenu("官方模型")
        # 添加子菜单项
        official_models = ["yolov5s", "yolov5m", "yolov5l", "yolov5x"]
        for model_name in official_models:
            action = QAction(model_name, self)
            action.triggered.connect(lambda checked, name=model_name: self.load_model(name))
            model_menu.addAction(action)

        # 添加“全屏”按钮
        fullscreen_action = QAction("全屏", self)
        fullscreen_action.triggered.connect(self.toggle_fullscreen)

        # 将 QAction 转换为 QToolButton
        fullscreen_button = QToolButton(self.menu_bar)
        fullscreen_button.setDefaultAction(fullscreen_action)  # 将 QAction 绑定到 QToolButton
        fullscreen_button.setFixedSize(300, 50)  # 设置固定大小为 100x50 像素
        fullscreen_button.setStyleSheet("""
            QToolButton {
                padding: 5px;  /* 设置内边距 */
                margin-right: 100px;  /* 设置右边距 */
            }
        """)  # 移除按钮边框

        fullscreen_button.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Expanding)  # 宽度固定，高度自适应
        
        # 调整 QMenuBar 的样式，为 QToolButton 预留空间
        self.menu_bar.setStyleSheet("""
            QMenuBar {
                padding-right: 120px;  /* 为 QToolButton 预留空间 */
            }
        """)

        # 将“全屏”按钮放置在菜单栏的最右侧
        self.menu_bar.setCornerWidget(fullscreen_button, Qt.TopRightCorner)  # 使用 QToolButton 作为 corner widget

        # 设置菜单栏样式
        self.menu_bar.setStyleSheet("""
            QMenuBar {
                background-color: #1E222A;  /* 更深的背景色 */
                color: #D8DEE9;  /* 更亮的文字颜色 */
                padding: 10px;
            }
            QMenuBar::item {
                background-color: transparent;
                padding: 5px 10px;
            }
            QMenuBar::item:selected {
                background-color: #3B4252;  /* 更深的选中背景色 */
            }
            QMenu {
                background-color: #2E3440;  /* 更深的菜单背景色 */
                color: #D8DEE9;  /* 更亮的文字颜色 */
                border: 1px solid #4C566A;
            }
            QMenu::item {
                padding: 5px 20px;
            }
            QMenu::item:selected {
                background-color: #4C566A;  /* 更深的选中背景色 */
            }
            QToolButton {
       
            padding: 5px;
            margin-right: 20px;  /* 进一步增加右边距 */
        }

        """)

        # 将按钮添加到布局中
        button_panel.addWidget(self.btn_start)
        button_panel.addWidget(self.btn_stop)
        button_panel.addWidget(self.btn_clear_roi)
        button_panel.addWidget(self.btn_draw_roi)

        right_layout.addLayout(button_panel)

        # 设置右侧区域的最小宽度
        right_panel.setMinimumWidth(1600)  # 右侧区域最小宽度为1600像素
        right_panel.setMinimumHeight(1600)

        # 将左右面板添加到 QSplitter
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setStretchFactor(0, 1)  # 左侧占1份
        splitter.setStretchFactor(1, 2)  # 右侧占2份


        # 设置全局样式表
        self.setStyleSheet("""
            QMainWindow {
                background-color: #2E3440;
            }
            QGroupBox {
                background-color: #3B4252;
                color: #ECEFF4;
                border: 1px solid #4C566A;
                border-radius: 15px;
                margin-top: 10px;
                padding-top: 10px;
            }
            QLabel {
                color: #ECEFF4;
            }
            QPushButton {
                background-color: #4C566A;
                color: #ECEFF4;
                border: 1px solid #4C566A;
                border-radius: 15px;
                padding: 5px;
                font-size: 35px;
            }
            QPushButton:hover {
                background-color: #5E81AC;
            }
            QPushButton:pressed {
                background-color: #81A1C1;
            }
            QComboBox, QLineEdit, QSlider, QListWidget, QTableWidget {
                background-color: #4C566A;
                color: #ECEFF4;
                border: 1px solid #4C566A;
                border-radius: 15px;
                padding: 5px;
            }
            QTableWidget {
                gridline-color: #4C566A;
            }
            QHeaderView:section {
                background-color: #4C566A;
                color: #ECEFF4;
                padding: 5px;
                border: 1px solid #4C566A;
            }
            QScrollBar:vertical {
                background-color: #4C566A;
                width: 15px;
                margin: 15px 0 15px 0;
            }
            QScrollBar:handle:vertical {
                background-color: #81A1C1;
                min-height: 20px;
                border-radius: 15px;
            }
            QScrollBar:add-line:vertical, QScrollBar:sub-line:vertical {
                background: none;
            }
            QScrollBar::add-page:vertical, QScrollBar::sub-page:vertical {
                background: none;
            }
            QMessageBox {
                background-color: #2E3440;
            }
            QMessageBox QLabel {
                color: #ECEFF4;
            }
            QMessageBox QPushButton {
                background-color: #4C566A;
                color: #ECEFF4;
                border: 1px solid #4C566A;
                border-radius: 5px;
                padding: 5px;
            }
            QMessageBox QPushButton:hover {
                background-color: #5E81AC;
            }
        """)
            
    def select_all_classes(self):
        """全选所有类别"""
        for i in range(self.class_list.count()):
            item = self.class_list.item(i)
            item.setSelected(True)

    def deselect_all_classes(self):
        """取消所有选择"""
        for i in range(self.class_list.count()):
            item = self.class_list.item(i)
            item.setSelected(False)

    def load_local_model(self):
        """加载本地模型"""
        model_path, _ = QFileDialog.getOpenFileName(self, "选择模型文件", "", "模型文件 (*.pt)")
        if model_path:
            try:
                self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
                if torch.cuda.is_available():
                    self.model = self.model.cuda()
                self.statusBar().showMessage(f"本地模型 {model_path} 加载成功")

                # 从模型中获取类别列表
                if hasattr(self.model, 'names'):
                    if isinstance(self.model.names, dict):  # 如果 names 是字典
                        self.available_classes = list(self.model.names.values())  # 提取类别名称
                    else:  # 如果 names 是列表
                        self.available_classes = self.model.names
                else:  # 如果模型没有 names 属性
                    self.available_classes = self.get_coco_classes()
                
                self.update_class_list()  # 更新类别选择列表

                # 更新模型路径显示和工具提示
                # self.btn_load_model.setText(f"模型路径: {model_path[:20]}...")  # 限制显示长度
                self.btn_load_model.setToolTip(model_path)  # 显示完整路径
                if self.camera_thread:
                    self.camera_thread.model = self.model
            except Exception as e:
                QMessageBox.critical(self, "错误", f"模型加载失败: {str(e)}")

    def toggle_draw_roi(self):
        """切换绘制 ROI 的状态"""
        if not self.camera_thread or not self.camera_thread.isRunning():
            QMessageBox.warning(self, "警告", "请先开始检测，再绘制 ROI")
            return

        self.is_drawing_roi = not self.is_drawing_roi
        if self.is_drawing_roi:
            self.btn_draw_roi.setText("停止绘制 ROI")
            self.points = []  # 重置顶点
        else:
            self.btn_draw_roi.setText("开始绘制 ROI")
            self.update_roi_mask()  # 完成绘制并生成掩码

    def update_iou_thres_from_slider(self):
        iou_thres = self.iou_slider.value() / 100
        self.iou_input.setText(f"{iou_thres:.2f}")
        if self.camera_thread:
            self.camera_thread.iou_thres = iou_thres

    def update_iou_thres_from_input(self):
        try:
            iou_thres = float(self.iou_input.text())
            if 0 <= iou_thres <= 1:
                self.iou_slider.setValue(int(iou_thres * 100))
                if self.camera_thread:
                    self.camera_thread.iou_thres = iou_thres
            else:
                QMessageBox.warning(self, "警告", "IoU阈值必须在 0 到 1 之间")
        except ValueError:
            QMessageBox.warning(self, "警告", "请输入有效的数字")

    def mouse_press_event(self, event):
        """鼠标点击事件：记录多边形顶点"""
        if self.is_drawing_roi and event.button() == Qt.LeftButton:
            self.drawing = True
            # 获取点击位置的坐标
            x, y = event.x(), event.y()
            self.points.append((x, y))
            self.update_video_label()

    def mouse_move_event(self, event):
        """鼠标移动事件：实时更新多边形轮廓"""
        if self.is_drawing_roi and self.drawing:
            self.update_video_label()

    def mouse_release_event(self, event):
        """鼠标释放事件：结束当前绘制"""
        if self.is_drawing_roi and event.button() == Qt.LeftButton:
            self.drawing = False

    def update_video_label(self):
        """更新视频显示区域，实时绘制多边形轮廓"""
        if hasattr(self, 'camera_thread') and self.camera_thread and hasattr(self.camera_thread, 'current_frame'):
            frame = self.camera_thread.current_frame.copy()
            if len(self.points) >= 2:
                # 绘制多边形轮廓
                pts = np.array(self.points, np.int32)
                pts = pts.reshape((-1, 1, 2))
                cv2.polylines(frame, [pts], isClosed=False, color=(0, 255, 0), thickness=2)
                # 绘制当前点到上一个点的连线
                if len(self.points) > 1:
                    cv2.line(frame, self.points[-2], self.points[-1], (0, 255, 0), 2)
            # 转换图像格式
            rgb_image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w
            qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
            self.video_label.setPixmap(QPixmap.fromImage(qt_image))

    def update_roi_mask(self):
        """生成 ROI 掩码"""
        if len(self.points) >= 3:  # 至少需要3个点才能形成多边形
            # 创建掩码
            mask = np.zeros((self.video_label.height(), self.video_label.width()), dtype=np.uint8)
            pts = np.array(self.points, np.int32)
            pts = pts.reshape((-1, 1, 2))
            cv2.fillPoly(mask, [pts], 255)
            self.roi_mask = mask
            if self.camera_thread:
                self.camera_thread.roi_mask = self.roi_mask
            # 显示 ROI 区域
            self.update_video_label()
        else:
            QMessageBox.warning(self, "警告", "至少需要3个点才能形成多边形")

    def clear_roi(self):
        """清除 ROI 区域"""
        self.roi_mask = None
        self.points = []
        if self.camera_thread:
            self.camera_thread.roi_mask = None
        # 更新视频显示区域
        self.update_video_label()

    def load_model(self, model_name='yolov5s'):
        """加载 YOLOv5 模型，如果本地没有则自动下载"""
        model_path = os.path.join(self.model_dir, f"{model_name}.pt")
        try:
            self.model = torch.hub.load('ultralytics/yolov5', model_name, pretrained=True)
            # 保存模型到本地
            torch.save(self.model.state_dict(), model_path)
            self.statusBar().showMessage(f"模型 {model_name} 已下载并保存到 {model_path}")
            if torch.cuda.is_available():
                self.model = self.model.cuda()
            self.statusBar().showMessage(f"模型 {model_name} 加载成功")

            # 从模型中获取类别列表
            if hasattr(self.model, 'names'):
                if isinstance(self.model.names, dict):  # 如果 names 是字典
                    self.available_classes = list(self.model.names.values())  # 提取类别名称
                else:  # 如果 names 是列表
                    self.available_classes = self.model.names
            else:  # 如果模型没有 names 属性
                self.available_classes = self.get_coco_classes()
            
            self.update_class_list()  # 更新类别选择列表            
        
        except Exception as e:
            QMessageBox.critical(self, "错误", f"模型加载失败: {str(e)}")
            sys.exit(1)

    def change_model(self):
        model_name = self.model_selector.currentText()
        self.load_model(model_name)
        if self.camera_thread:
            self.camera_thread.model = self.model

    def update_class_list(self):
        """更新类别选择列表"""
        self.class_list.clear()  # 清空现有列表
        for cls in self.available_classes:
            self.class_list.addItem(cls)

    def get_coco_classes(self):
        # YOLOv5 支持的 COCO 数据集类别
        return [
            'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
            'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
            'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
            'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
            'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
            'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
            'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
            'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
            'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
            'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
            'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
            'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
            'scissors', 'teddy bear', 'hair drier', 'toothbrush'
        ]

    def update_ip_camera_input_completer(self):
        """更新 IP 摄像头输入框的下拉菜单"""
        if hasattr(self, 'ip_camera_input'):  # 确保 ip_camera_input 已定义
            completer = QCompleter(self.history_ip_addresses, self)
            completer.setCaseSensitivity(Qt.CaseInsensitive)  # 不区分大小写
            completer.setFilterMode(Qt.MatchContains)  # 允许部分匹配
            self.ip_camera_input.setCompleter(completer)

    def update_source_ui(self):
        source_type = self.source_selector.currentText()
        if source_type == "IP摄像头":
            self.ip_camera_label.setVisible(True)
            self.ip_camera_input.setVisible(True)
            self.file_button.setVisible(False)
            self.btn_clear_history.setVisible(True)  # 显示清除历史按钮

            # 当用户输入 IP 地址并按下回车时，保存历史记录
            self.ip_camera_input.returnPressed.connect(self.save_current_ip_address)

        elif source_type in ["视频文件", "图片文件"]:
            self.ip_camera_label.setVisible(False)
            self.ip_camera_input.setVisible(False)
            self.file_button.setVisible(True)
            self.btn_clear_history.setVisible(False)  # 隐藏清除历史按钮
        else:  # 本地摄像头
            self.ip_camera_label.setVisible(False)
            self.ip_camera_input.setVisible(False)
            self.file_button.setVisible(False)
            self.btn_clear_history.setVisible(False)  # 隐藏清除历史按钮

    def save_current_ip_address(self):
        """保存当前输入的 IP 地址到历史记录"""
        current_ip = self.ip_camera_input.text().strip()
        if current_ip and current_ip not in self.history_ip_addresses:
            self.history_ip_addresses.append(current_ip)
            self.save_history_ip_addresses()  # 保存到文件
            self.update_ip_camera_input_completer()  # 更新下拉菜单
            print("保存的 IP 地址:", current_ip)  # 调试信息

    def select_file(self):
        source_type = self.source_selector.currentText()
        if source_type == "视频文件":
            file_path, _ = QFileDialog.getOpenFileName(self, "选择视频文件", "", "视频文件 (*.mp4 *.avi *.mov)")
        elif source_type == "图片文件":
            file_path, _ = QFileDialog.getOpenFileName(self, "选择图片文件", "", "图片文件 (*.jpg *.png *.bmp)")
        if file_path:
            self.source_path = file_path
            if source_type == "图片文件":
                self.load_and_display_image(file_path)  # 选择图片后立即显示
            self.statusBar().showMessage(f"已选择文件: {self.source_path}")    
            print(f"已选择文件: {self.source_path}")  # 调试信息

        else:
            QMessageBox.warning(self, "警告", "未选择文件")

    def load_and_display_image(self, file_path):
        """加载图片并调整大小以适应检测区域"""
        image = cv2.imread(file_path)
        if image is None:
            QMessageBox.warning(self, "警告", "无法加载图片")
            return

        # 调整图片大小以适应检测区域
        height, width, _ = image.shape
        label_width = self.video_label.width()
        label_height = self.video_label.height()
        scale = min(label_width / width, label_height / height)
        new_width = int(width * scale)
        new_height = int(height * scale)
        resized_image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
        
        # 将图片显示在检测区域
        rgb_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        qt_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        self.video_label.setPixmap(QPixmap.fromImage(qt_image))

    def select_save_path(self):
        """选择保存路径，并生成唯一的文件名"""
        # 获取当前时间戳
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        # 默认文件名
        default_file_name = f"{timestamp}.mp4"
        
        # 打开文件对话框选择保存路径
        file_path, _ = QFileDialog.getSaveFileName(
            self, 
            "选择保存路径", 
            default_file_name,  # 默认文件名
            "MP4文件 (*.mp4)"
        )

        if file_path:
            self.output_video_path = file_path
            # 更新保存路径显示和工具提示
            self.btn_save_path.setText(f"保存路径: {file_path[:20]}...")  # 限制显示长度
            self.btn_save_path.setToolTip(file_path)  # 显示完整路径
            if self.camera_thread:
                self.camera_thread.output_video_path = self.output_video_path

    def on_camera_thread_finished(self):
        """当检测线程完成时调用"""
        if self.source_selector.currentText() == "图片文件":
            self.stop_camera()

    def on_detection_finished(self):
        """当检测完成时调用"""
        
        if self.source_selector.currentText() == "图片文件":
            self.statusBar().showMessage(f"检测成功！！！" )# 调试信息
            self.stop_camera()

    def start_camera(self):
        if not self.camera_thread:
            # 获取多选的目标类别
            selected_classes = [item.text() for item in self.class_list.selectedItems()]
            source_type = self.source_selector.currentText()
            if source_type == "本地摄像头":
                source_path = 0  # 默认摄像头
                source_type = "camera"
            elif source_type == "IP摄像头":
                source_path = self.ip_camera_input.text()
                source_type = "ip_camera"
            elif source_type == "视频文件":
                if not hasattr(self, 'source_path') or self.source_path is None:
                    QMessageBox.warning(self, "警告", "请先选择视频文件")
                    return
                source_path = self.source_path
                source_type = "video"
            elif source_type == "图片文件":
                if not hasattr(self, 'source_path') or self.source_path is None:
                    QMessageBox.warning(self, "警告", "请先选择图片文件")
                    return
                source_path = self.source_path
                source_type = "image"
            # 检查文件路径是否存在
            if source_type in ["video", "image"] and not os.path.exists(source_path):
                QMessageBox.warning(self, "警告", f"文件路径无效: {source_path}")
                return    

            # 获取 video_label 的宽度和高度
            target_width = self.video_label.width()
            target_height = self.video_label.height()

            # 生成新的视频文件名，避免覆盖
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            self.output_video_path = f"./runs/videos/{timestamp}.mp4"

            self.camera_thread = CameraThread(
                self.model,
                conf_thres=self.conf_slider.value() / 100,
                iou_thres=self.iou_slider.value() / 100,  # 传递IoU阈值
                selected_classes=selected_classes,
                source_type=source_type,
                source_path=source_path,
                roi_mask=self.roi_mask,  # 传递ROI掩码
                output_video_path=self.output_video_path,  # 传递保存路径
                target_width=target_width,  # 传递目标宽度
                target_height=target_height,  # 传递目标高度
                save_dir=self.save_dir  # 传递保存目录

            )
            self.camera_thread.update_frame.connect(self.update_frame)
            self.camera_thread.update_fps.connect(self.update_fps)
            self.camera_thread.detection_finished.connect(self.on_detection_finished)  # 连接检测完成信号
            self.camera_thread.start()
            
            self.btn_start.setEnabled(False)
            self.btn_stop.setEnabled(True)
            self.btn_draw_roi.setEnabled(True)  # 开始检测后禁用绘制 ROI 按钮
            self.btn_clear_roi.setEnabled(True)  # 开始检测后禁用清除 ROI 按钮

    def stop_camera(self):
        if self.camera_thread:
            self.camera_thread.stop()
            self.camera_thread = None
            self.btn_start.setEnabled(True)
            self.btn_stop.setEnabled(False)
            self.btn_draw_roi.setEnabled(True)  # 停止检测后启用绘制 ROI 按钮
            self.btn_clear_roi.setEnabled(True)  # 停止检测后启用清除 ROI 按钮

                # 仅在检测视频或摄像头时清除检测区域的图片
            if self.source_selector.currentText() in ["本地摄像头", "IP摄像头", "视频文件"]:
                self.video_label.clear()
                self.video_label.setAlignment(Qt.AlignCenter)
            else:
            # 如果是图片文件，可以在这里添加其他逻辑
                self.statusBar().showMessage("图片检测成功！！！")    
            
    def update_selected_classes(self):
        """更新目标类别"""
        if self.camera_thread:
            selected_classes = [item.text() for item in self.class_list.selectedItems()]
            self.camera_thread.update_selected_classes(selected_classes)

    def update_frame(self, image, detection_list):
        """更新视频帧"""
        if not self.camera_thread or not self.camera_thread.isRunning():
                return  # 如果线程已停止，不更新画面
        # 将 QImage 转换为 QPixmap
        pixmap = QPixmap.fromImage(image)
        # 动态调整图像大小以适应窗口
        self.video_label.setPixmap(
            pixmap.scaled(self.video_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
        
        # 更新检测结果表格
        self.update_results_table(detection_list)

    def update_results_table(self, detection_list):
        # 统计每个类别的数量和置信度
        class_stats = {}
        for detection in detection_list:
            cls = detection["class"]
            conf = float(detection["confidence"])
            if cls not in class_stats:
                class_stats[cls] = {"count": 0, "confidence": []}
            class_stats[cls]["count"] += 1
            class_stats[cls]["confidence"].append(conf)

        # 更新表格
        self.results_table.setRowCount(len(class_stats))
        for i, (cls, stats) in enumerate(class_stats.items()):
            avg_conf = sum(stats["confidence"]) / len(stats["confidence"])

            # 创建 QTableWidgetItem 并设置居中对齐
            class_item = QTableWidgetItem(cls)
            class_item.setTextAlignment(Qt.AlignCenter)  # 居中对齐
            self.results_table.setItem(i, 0, class_item)

            conf_item = QTableWidgetItem(f"{avg_conf:.2f}")
            conf_item.setTextAlignment(Qt.AlignCenter)  # 居中对齐
            self.results_table.setItem(i, 1, conf_item)

            count_item = QTableWidgetItem(str(stats["count"]))
            count_item.setTextAlignment(Qt.AlignCenter)  # 居中对齐
            self.results_table.setItem(i, 2, count_item)

          # 更新历史数据
        current_time = self.start_time.msecsTo(QDateTime.currentDateTime()) / 1000  # 当前时间（秒）
        for cls, stats in class_stats.items():
            if cls not in self.history_data:
                self.history_data[cls] = {"time": [], "count": []}
            self.history_data[cls]["time"].append(current_time)
            self.history_data[cls]["count"].append(stats["count"])    

    def save_history_data(self):
        """每隔 3 秒保存一次历史数据"""
        current_time = self.start_time.msecsTo(QDateTime.currentDateTime()) / 1000  # 当前时间（秒）
        for cls in self.history_data:
            if cls not in self.saved_history_data:
                self.saved_history_data[cls] = {"time": [], "count": []}
            # 保存最近 1 秒的数据
            for i, t in enumerate(self.history_data[cls]["time"]):
                if current_time - t <= 1:  # 只保存最近 1 秒的数据
                    self.saved_history_data[cls]["time"].append(t)
                    self.saved_history_data[cls]["count"].append(self.history_data[cls]["count"][i])
        # 清空当前历史数据
        self.history_data = {}

    def update_fps(self, fps):
        self.fps_label.setText(f"FPS: {fps:.1f}")

    def update_conf_thres_from_slider(self):
        conf_thres = self.conf_slider.value() / 100
        self.conf_input.setText(f"{conf_thres:.2f}")
        if self.camera_thread:
            self.camera_thread.conf_thres = conf_thres

    def update_conf_thres_from_input(self):
        try:
            conf_thres = float(self.conf_input.text())
            if 0 <= conf_thres <= 1:
                self.conf_slider.setValue(int(conf_thres * 100))
                if self.camera_thread:
                    self.camera_thread.conf_thres = conf_thres
            else:
                QMessageBox.warning(self, "警告", "置信度必须在 0 到 1 之间")
        except ValueError:
            QMessageBox.warning(self, "警告", "请输入有效的数字")

    def resizeEvent(self, event):
        # 窗口大小变化时动态调整图像
        if self.camera_thread:
            self.update_frame(self.video_label.pixmap().toImage(), [])
        super().resizeEvent(event)

    def closeEvent(self, event):
        self.stop_camera()
        super().closeEvent(event)

    def update_category_combo(self, combo_box):
        """更新类别选择框"""
        combo_box.clear()  # 清空选择框
        for i in range(self.class_list.count()):  # 遍历目标类别选择列表
            item = self.class_list.item(i)
            if item.isSelected():  # 只加载选中的类别
                combo_box.addItem(item.text())

    def update_selected_classes(self):
        """更新目标类别"""
        if self.camera_thread:
            selected_classes = [item.text() for item in self.class_list.selectedItems()]
            self.camera_thread.update_selected_classes(selected_classes)
        
        # 更新折线图的类别选择框
        if hasattr(self, 'line_chart_dialog') and self.line_chart_dialog.isVisible():
            self.update_category_combo(self.category_combo)            

    def show_line_chart(self):
        """显示折线图"""
        self.line_chart_dialog = QDialog(self)
        self.line_chart_dialog.setWindowTitle("折线图")
        self.line_chart_dialog.setGeometry(800, 100, 1200, 1200)

        layout = QVBoxLayout(self.line_chart_dialog)

        # 添加类别选择框
        self.category_combo = QComboBox(self.line_chart_dialog)
        self.update_category_combo(self.category_combo)  # 传递 self.category_combo 作为参数  # 初始化加载目标类别选择列表
        layout.addWidget(self.category_combo)


        # 监听类别选择框的变化
        self.category_combo.currentTextChanged.connect(self.update_line_chart)

        # 添加画布
        self.figure = Figure()
        self.canvas = FigureCanvas(self.figure)
        layout.addWidget(self.canvas)

        self.line_chart_dialog.setLayout(layout)
        self.line_chart_dialog.show()

        # 重置历史数据和计时器
        self.history_data = {}  # 清空当前历史数据
        self.saved_history_data = {}  # 清空保存的历史数据
        self.start_time = QDateTime.currentDateTime()  # 重置启动时间

        # 启动定时器，每隔 1 秒更新一次折线图
        self.timer.start(1000)
        
        

        # 初始绘制
        self.update_line_chart()

    def update_line_chart(self):
        """更新折线图"""
        if not hasattr(self, 'line_chart_dialog') or not self.line_chart_dialog.isVisible():
            return  # 如果折线图对话框未打开，则不更新
        
        category = self.category_combo.currentText()

        # 获取当前时间
        current_time = self.start_time.msecsTo(QDateTime.currentDateTime()) / 1000  # 当前时间（秒）

        # 过滤掉超过 3 秒的数据
        if category in self.history_data:
            x = [
                t for t in self.history_data[category]["time"]
                if current_time - t <= 1  # 只保留最近 3 秒的数据
            ]
            y = [
                c for i, c in enumerate(self.history_data[category]["count"])
                if current_time - self.history_data[category]["time"][i] <= 1  # 只保留最近 3 秒的数据
            ]
        else:
            x, y = [], []

        # 清空画布
        self.figure.clear()

        # 设置全局字体大小
        plt.rcParams.update({'font.size': 20})  # 设置字体大小为14

         # 绘制折线图
        ax = self.figure.add_subplot(111)
        ax.plot(x, y, marker='o', linestyle='-', color='b', label=category)
        ax.set_title(f"{category} 数量统计", fontsize=30)  # 设置标题字体大小
        ax.set_xlabel("时间 (秒)", fontsize=25)  # 设置x轴标签字体大小
        ax.set_ylabel("数量", fontsize=25)  # 设置y轴标签字体大小

        # 设置x轴刻度
        if x:
            ax.set_xticks(x)
            ax.set_xticklabels([f"{t:.1f}s" for t in x], rotation=45, ha='right')

        ax.set_ylim(0, 100)    

        # 添加图例
        ax.legend()    
        # 刷新画布
        self.canvas.draw()

    def load_history_ip_addresses(self):
        """从文件中加载历史 IP 摄像头地址"""
        if os.path.exists(self.history_file):
            try:
                with open(self.history_file, "r") as file:
                    self.history_ip_addresses = json.load(file)                                  
            except Exception as e:      
                self.history_ip_addresses = []  # 如果加载失败，初始化为空列表
        else:
            self.history_ip_addresses = []  # 如果文件不存在，初始化为空列表

    def save_history_ip_addresses(self):
        """将历史 IP 摄像头地址保存到文件"""
        try:
            with open(self.history_file, "w") as file:
                json.dump(self.history_ip_addresses, file)
                self.statusBar().showMessage(f"保存IP地址成功")# 调试信息
        except Exception as e:
            self.statusBar().showMessage(f"保存历史 IP 地址失败: {e}")# 调试信息
            

    def clear_history_ip_addresses(self):
        """清除历史 IP 地址"""
        # 弹出确认对话框
        reply = QMessageBox.question(
            self,
            "确认清除",
            "确定要清除所有历史 IP 地址吗？",
            QMessageBox.Yes | QMessageBox.No,
            QMessageBox.No
        )

        if reply == QMessageBox.Yes:
            # 清空历史记录
            self.history_ip_addresses = []
            self.save_history_ip_addresses()  # 保存空列表到文件
            self.update_ip_camera_input_completer()  # 更新自动补全
            QMessageBox.information(self, "成功", "历史 IP 地址已清除")
    
    def toggle_fullscreen(self):
        """切换全屏模式"""
        
        if self.isFullScreen():
            self.showNormal()  # 退出全屏模式
            self.menu_bar.cornerWidget().setText("全屏")  # 直接更新 QToolButton 的文本
            
        else:
            self.showFullScreen()  # 进入全屏模式
            self.menu_bar.cornerWidget().setText("退出全屏")  # 直接更新 QToolButton 的文本
            
        # 手动刷新 QToolButton
        self.menu_bar.cornerWidget().update()    


if __name__ == "__main__":
    app = QApplication(sys.argv)
    
    # 设置调色板以自定义颜色
    palette = QPalette()
    palette.setColor(QPalette.Window, QColor(53, 53, 53))  # 设置窗口背景色
    palette.setColor(QPalette.WindowText, QColor(255, 255, 255))  # 设置窗口文本颜色
    palette.setColor(QPalette.Base, QColor(35, 35, 35))  # 设置基础背景色
    palette.setColor(QPalette.AlternateBase, QColor(53, 53, 53))  # 设置交替背景色
    palette.setColor(QPalette.ToolTipBase, QColor(255, 255, 255))  # 设置工具提示背景色
    palette.setColor(QPalette.ToolTipText, QColor(255, 255, 255))  # 设置工具提示文本颜色
    palette.setColor(QPalette.Text, QColor(255, 255, 255))  # 设置文本颜色
    palette.setColor(QPalette.Button, QColor(53, 53, 53))  # 设置按钮背景色
    palette.setColor(QPalette.ButtonText, QColor(255, 255, 255))  # 设置按钮文本颜色
    palette.setColor(QPalette.BrightText, QColor(255, 0, 0))  # 设置高亮文本颜色
    palette.setColor(QPalette.Highlight, QColor(42, 130, 218))  # 设置高亮背景色
    palette.setColor(QPalette.HighlightedText, QColor(0, 0, 0))  # 设置高亮文本颜色

    app.setPalette(palette)
    # 设置应用程序样式为 Fusion
    app.setStyle("Fusion")
    window = YOLOv5CameraUI()
    window.show()
    sys.exit(app.exec_())